// Copyright (C) 2023 Computer Vision Lab, Electrical Engineering, 
// Indian Institute of Science, Bengaluru, India.
// All rights reserved.
// 
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//     * Redistributions of source code must retain the above copyright
//       notice, this list of conditions and the following disclaimer.
//     * Redistributions in binary form must reproduce the above
//       copyright notice, this list of conditions and the following
//       disclaimer in the documentation and/or other materials provided
//       with the distribution.
//     * Neither the name of Indian Institute of Science nor the
//       names of its contributors may be used to endorse or promote products
//       derived from this software without specific prior written permission.
// 
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS 
// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL 
// THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING 
// FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER 
// DEALINGS IN THE SOFTWARE.
// 
// 
// This file is a part of the implementation for the paper:
// Chitturi Sidhatha, Lalit Manam and Venu Madhav Govindu, 
// Adaptive Annealing for Robust Geometric Estimation,
// Proceedings of the Computer Vision and Pattern Recognition 2023
// 
// This code base is adapted from the code release of Fast Global Registration,
// Qian-Yi Zhou, Jaesik Park, and Vladlen Koltun, ECCV 2016

#include <vector>
#include <limits>
#include <flann/flann.hpp>

#include <numeric>      // std::iota
#include <algorithm>    // std::sort, std::stable_sort

#include <Eigen/Core>
#include <Eigen/Dense>
#include <Eigen/SVD>
#include <Eigen/LU> 
#include <Eigen/Geometry>
#include <unsupported/Eigen/CXX11/TensorSymmetry>

#define USE_ABSOLUTE_SCALE	0		// Measure distance in absolute scale (1) or in scale relative to the diameter of the model (0)
// #define MAX_CORR_DIST		0.025	// Maximum correspondence distance (also see comment of USE_ABSOLUTE_SCALE)
#define ITERATION_NUMBER	200		// Maximum number of iteration, changed by Sid from 64 to 200
// #define TUPLE_SCALE			0.95	// Similarity measure used for tuples of feature points.
// #define TUPLE_MAX_CNT		1000	// Maximum tuple numbers.
#define STOCH_SAMPLING 		false   // Stochastic sampling disabled by default.
#define SIG_MAX				1000    // Maximum value of sigma, or the starting sigma.
#define USE_UMEYAMA			true	// Should I use Umeyama or SE3Reg as the local optimizer
#define ANNEAL_FACTOR 		1.4     // Fixed annealing factor of 1.4 in case of fixed annealing
#define IS_FIXED_ANNEAL		false   // Fixed annealing scheme is disabled by default.
#define IS_PERTURB				true    // Flips are enabled by default.
#define EIG_MIN_THRES		0.0     // Minimum eigen value threshold for checking pdness of the Hessian
#define SIG_MIN 			0.1 	// Minimum sigma where the GNC should stop
#define THRES				0.01    // A threshold which HAS TO BE DEFINED LATER (NEEDS CHECKING)
// #define M_PI           3.14159265358979323846 

namespace gom {
  
typedef std::vector<Eigen::Vector3f> Points;
typedef std::vector<std::pair<int, int> > Correspondences;
typedef Eigen::MatrixXd PointsMat;
typedef Eigen::Vector<bool, Eigen::Dynamic> VectorXb;
typedef Eigen::Tensor<double,3> Tensor3d;
typedef Eigen::Array<double,5,1> Array5d;
typedef Eigen::Array<double,4,1> Array4d;
typedef Eigen::Array<double,3,1> Array3d;
typedef Eigen::Array<double,6,1> Array6d;

static void sort_indexes(const Eigen::VectorXd &v, Eigen::VectorXi& idx);
static void ProbSampleReplace(Eigen::VectorXd& p, Eigen::VectorXi& perm, int nans, Eigen::VectorXi& ans);
void generateGemanMcClureApprox1(Eigen::VectorXd& rjj,double sig,Eigen::MatrixXd& z1,Eigen::MatrixXd& z2);
double getMinEigVal(const Eigen::MatrixXd& H);

Eigen::Matrix3d vec2crossmat(Eigen::Vector3d v);
void w2R(Eigen::Vector3d& w, Eigen::Matrix3d& R);
Eigen::Matrix3d getRandomRotation();

class CApp{
public:
	CApp(bool 	stoch_sampling     = STOCH_SAMPLING,
		double 	sig_max 	       = SIG_MAX,
		bool 	isumeya			   = USE_UMEYAMA,
		double   anneal_factor      = ANNEAL_FACTOR,
		bool 	fixed_anneal 	   = IS_FIXED_ANNEAL,
		bool 	isperturb			   = IS_PERTURB,
		double 	eig_min_thres	   = EIG_MIN_THRES,
		double 	sig_min			   = SIG_MIN,
		double 	thres			   = THRES):
		stoch_sampling_(stoch_sampling),
		sig_max_(sig_max),
		isumeya_(isumeya),
		anneal_factor_(anneal_factor),
		fixed_anneal_(fixed_anneal),
		isperturb_(isperturb),
		eig_min_thres_(eig_min_thres),
		sig_min_(sig_min),
		thres_(thres){}
	void LoadPointCloud(const PointsMat& pts1, const PointsMat& pts2);
	void NormalizePoints();
	void SetMatches();
	Eigen::Matrix4d GetOutputTrans();
	Eigen::Matrix4d WeightedUmeyama(const Eigen::MatrixXd& pts1, const Eigen::MatrixXd& pts2, const Eigen::VectorXd& weights);
	double RobustUmeyamaPairwise();
	void GetTensorValue();
	int n_runs_ = 1;
	int n_times_inner_ = 0;
	double sig_stop_ = 0.1d;
	int n_iter_ = 0;
	int n_perturbs_ = 0;
	double t_gom_ = 0.0d;
private:
	// containers
	std::vector<Points> pointcloud_;
	PointsMat ptsA_, ptsB_;
	Eigen::Matrix4d TransOutput_;
	std::vector<std::pair<int, int> > corres_;

	// for normalization
	Points Means;
	double GlobalScale = 1.0f;
	double StartScale = 1.0f;

	// some internal functions
	void TransformPoints(Points& points, const Eigen::Matrix4f& Trans);
	// template <typename T> 
	// void sort_indexes(const std::vector<T> &v, const std::vector<std::size_t>& idx); 	
	int max_iters_ = ITERATION_NUMBER;
	bool 	stoch_sampling_;
	double 	sig_max_;
	bool 	isumeya_;
	double   anneal_factor_;
	bool 	fixed_anneal_;
	bool 	isperturb_;
	double 	eig_min_thres_;
	double 	sig_min_;
	double 	thres_;
	int 	num_corr_;


};

}
